import os
import time
import math
import random
import urllib.request
import socket
import shutil
from urllib.error import URLError

def clean_empty_folders(base_dir="tiles"):
    """清理空文件夹"""
    print("\nCleaning empty folders...")
    cleaned = 0
    for root, dirs, files in os.walk(base_dir, topdown=False):
        for dir_name in dirs:
            dir_path = os.path.join(root, dir_name)
            try:
                # 如果文件夹为空，删除它
                if not os.listdir(dir_path):
                    os.rmdir(dir_path)
                    cleaned += 1
                    print(f"Removed empty folder: {dir_path}")
            except Exception as e:
                print(f"Error cleaning folder {dir_path}: {str(e)}")
    print(f"Cleaned {cleaned} empty folders")

def check_image_valid(file_path):
    """检查图片文件是否有效"""
    try:
        if not os.path.exists(file_path):
            return False
        
        # 检查文件大小是否太小（可能是错误的图片）
        if os.path.getsize(file_path) < 1000:  # 1KB
            print(f"Found invalid file (too small): {file_path}")
            os.remove(file_path)  # 删除无效文件
            return False
        
        return True
    except Exception as e:
        print(f"Error checking file {file_path}: {str(e)}")
        return False

def download_tile(z, x, y, retry_count=5):
    """下载单个瓦片，支持重试"""
    file_path = f"tiles/{z}/{x}/{y}.png"
    
    # 确保目录存在
    os.makedirs(f"tiles/{z}/{x}", exist_ok=True)
    
    # ArcGIS瓦片URL
    url = f"https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}"
    
    # 设置请求头
    headers = {
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
        'Accept': 'image/png',
        'Accept-Language': 'en-US,en;q=0.5',
        'Accept-Encoding': 'gzip, deflate',
        'Connection': 'keep-alive',
        'Referer': 'https://www.arcgis.com/'
    }

    for attempt in range(retry_count):
        try:
            # 设置60秒超时
            socket.setdefaulttimeout(60)
            
            req = urllib.request.Request(url, headers=headers)
            with urllib.request.urlopen(req) as response:
                if response.status == 200:
                    # 先下载到临时文件
                    temp_file = file_path + '.tmp'
                    with open(temp_file, 'wb') as f:
                        f.write(response.read())
                    
                    # 检查临时文件是否有效
                    if os.path.exists(temp_file) and os.path.getsize(temp_file) > 1000:
                        # 如果有效，移动到正确的位置
                        os.replace(temp_file, file_path)
                        print(f'Successfully downloaded: {z}/{x}/{y}')
                        time.sleep(random.uniform(0.5, 1.0))
                        return True
                    else:
                        print(f'Downloaded file invalid: {z}/{x}/{y}')
                        if os.path.exists(temp_file):
                            os.remove(temp_file)
                else:
                    print(f'Server returned status {response.status} for {z}/{x}/{y}')
        except (URLError, socket.timeout, ConnectionError) as e:
            print(f'Error downloading {z}/{x}/{y} (attempt {attempt + 1}/{retry_count}): {str(e)}')
            if attempt < retry_count - 1:
                wait_time = random.uniform(3 + attempt * 3, 7 + attempt * 3)
                print(f'Waiting {wait_time:.1f} seconds before retry...')
                time.sleep(wait_time)
            continue
        except Exception as e:
            print(f'Unexpected error for {z}/{x}/{y}: {str(e)}')
            if attempt < retry_count - 1:
                time.sleep(random.uniform(4, 8))
            continue
        finally:
            # 确保清理临时文件
            if os.path.exists(file_path + '.tmp'):
                os.remove(file_path + '.tmp')
    return False

def check_and_download_missing():
    """检查并下载缺失的瓦片"""
    # 千岛湖区域的瓦片范围
    z = 12
    x_range = range(3420, 3435)  # x瓦片范围
    y_range = range(1680, 1695)  # y瓦片范围
    
    total_tiles = len(x_range) * len(y_range)
    missing_tiles = []
    
    # 首先检查所有缺失或无效的瓦片
    print("Checking for missing or invalid tiles...")
    for x in x_range:
        for y in y_range:
            file_path = f"tiles/{z}/{x}/{y}.png"
            if not check_image_valid(file_path):
                missing_tiles.append((z, x, y))
                print(f"Found missing tile: {z}/{x}/{y}")
    
    print(f"\nFound {len(missing_tiles)} missing or invalid tiles out of {total_tiles} total tiles")
    
    if not missing_tiles:
        print("No missing tiles found!")
        return
    
    # 下载缺失的瓦片
    print("\nStarting to download missing tiles...")
    success_count = 0
    failed_tiles = []
    
    for i, (z, x, y) in enumerate(missing_tiles, 1):
        print(f"\nProcessing tile {i}/{len(missing_tiles)}: {z}/{x}/{y}")
        if download_tile(z, x, y):
            success_count += 1
        else:
            failed_tiles.append((z, x, y))
        
        # 显示进度
        print(f'Progress: {success_count}/{len(missing_tiles)} missing tiles downloaded')
        
        # 每下载5个瓦片后多等待一会
        if i % 5 == 0:
            wait_time = random.uniform(3, 6)
            print(f'Taking a short break for {wait_time:.1f} seconds...')
            time.sleep(wait_time)
    
    # 清理空文件夹
    clean_empty_folders()
    
    # 最终报告
    print(f"\nDownload completed!")
    print(f"Successfully downloaded {success_count} out of {len(missing_tiles)} missing tiles")
    if failed_tiles:
        print(f"Failed to download {len(failed_tiles)} tiles:")
        for z, x, y in failed_tiles:
            print(f"- Tile {z}/{x}/{y}")

if __name__ == '__main__':
    try:
        check_and_download_missing()
    except KeyboardInterrupt:
        print('\nDownload interrupted by user')
        # 清理空文件夹
        clean_empty_folders()
    except Exception as e:
        print(f'Error: {str(e)}')